#include "seq_mv.hpp"
#include "csr_matop.hpp"

template<typename idx_t, typename data_t, int dof>
seq_Vector<idx_t, data_t, dof>::seq_Vector(const idx_t _num) : num(_num), tot_len(_num * dof)
{
    data = new data_t [tot_len];
    this->operator=(0.0);// 初始化，避免分配内存后的数据内有nan
}

template<typename idx_t, typename data_t, int dof>
seq_Vector<idx_t, data_t, dof>::seq_Vector(const seq_Vector & model) : num(model.num), tot_len(model.tot_len)
{
    data = new data_t [tot_len];
    this->operator=(0.0);// 初始化，避免分配内存后的数据内有nan
}

template<typename idx_t, typename data_t, int dof>
seq_Vector<idx_t, data_t, dof>::~seq_Vector()
{
    if (data) delete [] data;
}

template<typename idx_t, typename data_t, int dof>
void seq_Vector<idx_t, data_t, dof>::operator=(data_t val)
{
    #pragma omp parallel for schedule(static)
    for (long long i = 0; i < tot_len; i++) {
        data[i] = val;
    }
}
#ifdef USE_FP16
template class seq_Vector<int, __fp16, 1>;
template class seq_Vector<int, __fp16, 2>;
template class seq_Vector<int, __fp16, 3>;
template class seq_Vector<int, __fp16, 4>;
#endif
template class seq_Vector<int, float , 1>;
template class seq_Vector<int, float , 2>;
template class seq_Vector<int, float , 3>;
template class seq_Vector<int, float , 4>;
template class seq_Vector<int, double, 1>;
template class seq_Vector<int, double, 2>;
template class seq_Vector<int, double, 3>;
template class seq_Vector<int, double, 4>;

template<typename idx_t, typename data_t, typename res_t, int dof>
res_t seq_vec_dot(const seq_Vector<idx_t, data_t, dof> & x, const seq_Vector<idx_t, data_t, dof> & y)
{
    assert(x.num == y.num);
    res_t dot = 0.0;

    #pragma omp parallel for reduction(+:dot) schedule(static)
    for (long long i = 0; i < x.tot_len; i++) {
        dot += (res_t) x.data[i] * (res_t) y.data[i];
    }

    return dot;
}
template double seq_vec_dot(const seq_Vector<int, float , 1> &, const seq_Vector<int, float , 1> &);
template double seq_vec_dot(const seq_Vector<int, float , 2> &, const seq_Vector<int, float , 2> &);
template double seq_vec_dot(const seq_Vector<int, float , 3> &, const seq_Vector<int, float , 3> &);
template double seq_vec_dot(const seq_Vector<int, float , 4> &, const seq_Vector<int, float , 4> &);
template double seq_vec_dot(const seq_Vector<int, double, 1> &, const seq_Vector<int, double, 1> &);
template double seq_vec_dot(const seq_Vector<int, double, 2> &, const seq_Vector<int, double, 2> &);
template double seq_vec_dot(const seq_Vector<int, double, 3> &, const seq_Vector<int, double, 3> &);
template double seq_vec_dot(const seq_Vector<int, double, 4> &, const seq_Vector<int, double, 4> &);

template<typename idx_t, typename data_t, typename res_t, int dof>
res_t seq_vec_inf_norm(const seq_Vector<idx_t, data_t, dof> & x, long long* max_index)
{
    res_t max_abs = 0.0;
    * max_index = -1;

    for (long long i = 0; i < x.tot_len; i++) {
        if (max_abs < fabs(x.data[i])) {
            max_abs = fabs(x.data[i]);
            * max_index = i;
        }
    }

    return max_abs;
}
template double seq_vec_inf_norm(const seq_Vector<int, float , 1> &, long long* );
template double seq_vec_inf_norm(const seq_Vector<int, float , 2> &, long long* );
template double seq_vec_inf_norm(const seq_Vector<int, float , 3> &, long long* );
template double seq_vec_inf_norm(const seq_Vector<int, float , 4> &, long long* );
template double seq_vec_inf_norm(const seq_Vector<int, double, 1> &, long long* );
template double seq_vec_inf_norm(const seq_Vector<int, double, 2> &, long long* );
template double seq_vec_inf_norm(const seq_Vector<int, double, 3> &, long long* );
template double seq_vec_inf_norm(const seq_Vector<int, double, 4> &, long long* );

template<typename idx_t, typename data_t, typename scalar_t, int dof>
void seq_vec_add(const seq_Vector<idx_t, data_t, dof> & v1, scalar_t alpha, 
                 const seq_Vector<idx_t, data_t, dof> & v2, seq_Vector<idx_t, data_t, dof> & v) 
{
    assert(v1.num == v2.num && v1.num == v.num);
    
    #pragma omp parallel for schedule(static)
    for (long long i = 0; i < v.tot_len; i++) {
        v.data[i] = v1.data[i] + alpha * v2.data[i];
    }
}
template void seq_vec_add(const seq_Vector_I32F32_DOF1 &, float , const seq_Vector_I32F32_DOF1 &, seq_Vector_I32F32_DOF1 &);
template void seq_vec_add(const seq_Vector_I32F32_DOF2 &, float , const seq_Vector_I32F32_DOF2 &, seq_Vector_I32F32_DOF2 &);
template void seq_vec_add(const seq_Vector_I32F32_DOF3 &, float , const seq_Vector_I32F32_DOF3 &, seq_Vector_I32F32_DOF3 &);
template void seq_vec_add(const seq_Vector_I32F32_DOF4 &, float , const seq_Vector_I32F32_DOF4 &, seq_Vector_I32F32_DOF4 &);
template void seq_vec_add(const seq_Vector_I32F64_DOF1 &, double, const seq_Vector_I32F64_DOF1 &, seq_Vector_I32F64_DOF1 &);
template void seq_vec_add(const seq_Vector_I32F64_DOF2 &, double, const seq_Vector_I32F64_DOF2 &, seq_Vector_I32F64_DOF2 &);
template void seq_vec_add(const seq_Vector_I32F64_DOF3 &, double, const seq_Vector_I32F64_DOF3 &, seq_Vector_I32F64_DOF3 &);
template void seq_vec_add(const seq_Vector_I32F64_DOF4 &, double, const seq_Vector_I32F64_DOF4 &, seq_Vector_I32F64_DOF4 &);
  
template<typename idx_t, typename data_t, int dof>
void seq_vec_copy(const seq_Vector<idx_t, data_t, dof> & src, seq_Vector<idx_t, data_t, dof> & dst)
{
    assert(src.num == dst.num);

    #pragma omp parallel for schedule(static)
    for (long long i = 0; i < src.tot_len; i++) {
        dst.data[i] = src.data[i];
    }
}
template void seq_vec_copy(const seq_Vector_I32F32_DOF1 &, seq_Vector_I32F32_DOF1 &);
template void seq_vec_copy(const seq_Vector_I32F32_DOF2 &, seq_Vector_I32F32_DOF2 &);
template void seq_vec_copy(const seq_Vector_I32F32_DOF3 &, seq_Vector_I32F32_DOF3 &);
template void seq_vec_copy(const seq_Vector_I32F32_DOF4 &, seq_Vector_I32F32_DOF4 &);
template void seq_vec_copy(const seq_Vector_I32F64_DOF1 &, seq_Vector_I32F64_DOF1 &);
template void seq_vec_copy(const seq_Vector_I32F64_DOF2 &, seq_Vector_I32F64_DOF2 &);
template void seq_vec_copy(const seq_Vector_I32F64_DOF3 &, seq_Vector_I32F64_DOF3 &);
template void seq_vec_copy(const seq_Vector_I32F64_DOF4 &, seq_Vector_I32F64_DOF4 &);

template<typename idx_t, typename data_t, typename scalar_t, int dof>
void seq_vec_mul_by_scalar(const scalar_t coeff, const seq_Vector<idx_t, data_t, dof> & src, seq_Vector<idx_t, data_t, dof> & dst) 
{
    assert(src.num == dst.num);

    #pragma omp parallel for schedule(static)
    for (long long i = 0; i < src.tot_len; i++) {
        dst.data[i] = coeff * src.data[i];
    }
}
template void seq_vec_mul_by_scalar(const float , const seq_Vector_I32F32_DOF1 &, seq_Vector_I32F32_DOF1 &);
template void seq_vec_mul_by_scalar(const float , const seq_Vector_I32F32_DOF2 &, seq_Vector_I32F32_DOF2 &);
template void seq_vec_mul_by_scalar(const float , const seq_Vector_I32F32_DOF3 &, seq_Vector_I32F32_DOF3 &);
template void seq_vec_mul_by_scalar(const float , const seq_Vector_I32F32_DOF4 &, seq_Vector_I32F32_DOF4 &);
template void seq_vec_mul_by_scalar(const double, const seq_Vector_I32F64_DOF1 &, seq_Vector_I32F64_DOF1 &);
template void seq_vec_mul_by_scalar(const double, const seq_Vector_I32F64_DOF2 &, seq_Vector_I32F64_DOF2 &);
template void seq_vec_mul_by_scalar(const double, const seq_Vector_I32F64_DOF3 &, seq_Vector_I32F64_DOF3 &);
template void seq_vec_mul_by_scalar(const double, const seq_Vector_I32F64_DOF4 &, seq_Vector_I32F64_DOF4 &);

template<typename idx_t, typename data_t, typename scalar_t, int dof>
void seq_vec_scale(const scalar_t coeff, seq_Vector<idx_t, data_t, dof> & vec)
{
    #pragma omp parallel for schedule(static)
    for (long long i = 0; i < vec.tot_len; i++) {
        vec.data[i] *= coeff;
    }
}
template void seq_vec_scale(const float , seq_Vector_I32F32_DOF1 &);
template void seq_vec_scale(const float , seq_Vector_I32F32_DOF2 &);
template void seq_vec_scale(const float , seq_Vector_I32F32_DOF3 &);
template void seq_vec_scale(const float , seq_Vector_I32F32_DOF4 &);
template void seq_vec_scale(const double, seq_Vector_I32F64_DOF1 &);
template void seq_vec_scale(const double, seq_Vector_I32F64_DOF2 &);
template void seq_vec_scale(const double, seq_Vector_I32F64_DOF3 &);
template void seq_vec_scale(const double, seq_Vector_I32F64_DOF4 &);

template<typename idx_t, typename data_t1, typename data_t2, int dof>
void seq_vec_elemwise_mul(seq_Vector<idx_t, data_t1, dof> & inout_vec, const seq_Vector<idx_t, data_t2, dof> & scaleplate)
{
    assert(inout_vec.num == scaleplate.num);

    #pragma omp parallel for schedule(static)
    for (long long i = 0; i < inout_vec.tot_len; i++) {
        inout_vec.data[i] *= scaleplate[i];
    }
}
extern template void seq_vec_elemwise_mul(seq_Vector_I32F32_DOF1 &, const seq_Vector_I32F32_DOF1 &);
extern template void seq_vec_elemwise_mul(seq_Vector_I32F32_DOF2 &, const seq_Vector_I32F32_DOF2 &);
extern template void seq_vec_elemwise_mul(seq_Vector_I32F32_DOF3 &, const seq_Vector_I32F32_DOF3 &);
extern template void seq_vec_elemwise_mul(seq_Vector_I32F32_DOF4 &, const seq_Vector_I32F32_DOF4 &);

extern template void seq_vec_elemwise_mul(seq_Vector_I32F64_DOF1 &, const seq_Vector_I32F64_DOF1 &);
extern template void seq_vec_elemwise_mul(seq_Vector_I32F64_DOF2 &, const seq_Vector_I32F64_DOF2 &);
extern template void seq_vec_elemwise_mul(seq_Vector_I32F64_DOF3 &, const seq_Vector_I32F64_DOF3 &);
extern template void seq_vec_elemwise_mul(seq_Vector_I32F64_DOF4 &, const seq_Vector_I32F64_DOF4 &);

extern template void seq_vec_elemwise_mul(seq_Vector_I32F32_DOF1 &, const seq_Vector_I32F64_DOF1 &);
extern template void seq_vec_elemwise_mul(seq_Vector_I32F32_DOF2 &, const seq_Vector_I32F64_DOF2 &);
extern template void seq_vec_elemwise_mul(seq_Vector_I32F32_DOF3 &, const seq_Vector_I32F64_DOF3 &);
extern template void seq_vec_elemwise_mul(seq_Vector_I32F32_DOF4 &, const seq_Vector_I32F64_DOF4 &);

extern template void seq_vec_elemwise_mul(seq_Vector_I32F64_DOF1 &, const seq_Vector_I32F32_DOF1 &);
extern template void seq_vec_elemwise_mul(seq_Vector_I32F64_DOF2 &, const seq_Vector_I32F32_DOF2 &);
extern template void seq_vec_elemwise_mul(seq_Vector_I32F64_DOF3 &, const seq_Vector_I32F32_DOF3 &);
extern template void seq_vec_elemwise_mul(seq_Vector_I32F64_DOF4 &, const seq_Vector_I32F32_DOF4 &);
#ifdef USE_FP16
extern template void seq_vec_elemwise_mul(seq_Vector_I32F64_DOF1 &, const seq_Vector_I32F16_DOF1 &);
extern template void seq_vec_elemwise_mul(seq_Vector_I32F64_DOF2 &, const seq_Vector_I32F16_DOF2 &);
extern template void seq_vec_elemwise_mul(seq_Vector_I32F64_DOF3 &, const seq_Vector_I32F16_DOF3 &);
extern template void seq_vec_elemwise_mul(seq_Vector_I32F64_DOF4 &, const seq_Vector_I32F16_DOF4 &);

extern template void seq_vec_elemwise_mul(seq_Vector_I32F32_DOF1 &, const seq_Vector_I32F16_DOF1 &);
extern template void seq_vec_elemwise_mul(seq_Vector_I32F32_DOF2 &, const seq_Vector_I32F16_DOF2 &);
extern template void seq_vec_elemwise_mul(seq_Vector_I32F32_DOF3 &, const seq_Vector_I32F16_DOF3 &);
extern template void seq_vec_elemwise_mul(seq_Vector_I32F32_DOF4 &, const seq_Vector_I32F16_DOF4 &);
#endif

template<typename idx_t, typename data_t1, typename data_t2, int dof>
void seq_vec_elemwise_div(seq_Vector<idx_t, data_t1, dof> & inout_vec, const seq_Vector<idx_t, data_t2, dof> & scaleplate)
{
    assert(inout_vec.num == scaleplate.num);

    #pragma omp parallel for schedule(static)
    for (long long i = 0; i < inout_vec.tot_len; i++) {
        inout_vec.data[i] /= scaleplate[i];
    }
}
extern template void seq_vec_elemwise_div(seq_Vector_I32F32_DOF1 &, const seq_Vector_I32F32_DOF1 &);
extern template void seq_vec_elemwise_div(seq_Vector_I32F32_DOF2 &, const seq_Vector_I32F32_DOF2 &);
extern template void seq_vec_elemwise_div(seq_Vector_I32F32_DOF3 &, const seq_Vector_I32F32_DOF3 &);
extern template void seq_vec_elemwise_div(seq_Vector_I32F32_DOF4 &, const seq_Vector_I32F32_DOF4 &);

extern template void seq_vec_elemwise_div(seq_Vector_I32F64_DOF1 &, const seq_Vector_I32F64_DOF1 &);
extern template void seq_vec_elemwise_div(seq_Vector_I32F64_DOF2 &, const seq_Vector_I32F64_DOF2 &);
extern template void seq_vec_elemwise_div(seq_Vector_I32F64_DOF3 &, const seq_Vector_I32F64_DOF3 &);
extern template void seq_vec_elemwise_div(seq_Vector_I32F64_DOF4 &, const seq_Vector_I32F64_DOF4 &);

extern template void seq_vec_elemwise_div(seq_Vector_I32F32_DOF1 &, const seq_Vector_I32F64_DOF1 &);
extern template void seq_vec_elemwise_div(seq_Vector_I32F32_DOF2 &, const seq_Vector_I32F64_DOF2 &);
extern template void seq_vec_elemwise_div(seq_Vector_I32F32_DOF3 &, const seq_Vector_I32F64_DOF3 &);
extern template void seq_vec_elemwise_div(seq_Vector_I32F32_DOF4 &, const seq_Vector_I32F64_DOF4 &);

extern template void seq_vec_elemwise_div(seq_Vector_I32F64_DOF1 &, const seq_Vector_I32F32_DOF1 &);
extern template void seq_vec_elemwise_div(seq_Vector_I32F64_DOF2 &, const seq_Vector_I32F32_DOF2 &);
extern template void seq_vec_elemwise_div(seq_Vector_I32F64_DOF3 &, const seq_Vector_I32F32_DOF3 &);
extern template void seq_vec_elemwise_div(seq_Vector_I32F64_DOF4 &, const seq_Vector_I32F32_DOF4 &);

#include "scal_kernels.hpp"

template<typename idx_t, typename data_t, typename calc_t, int dof>
CSRMatrix<idx_t, data_t, calc_t, dof>::CSRMatrix(): nrows(0), ncols(0), nnz(0)
{
}

template<typename idx_t, typename data_t, typename calc_t, int dof>
CSRMatrix<idx_t, data_t, calc_t, dof>::CSRMatrix(const idx_t _nrows, const idx_t _ncols): nrows(_nrows), ncols(_ncols)
{
    row_ptr = new idx_t [nrows + 1];
    
    #pragma omp parallel for schedule(static)
    for (idx_t i = 0; i < nrows + 1; i++)
        row_ptr[i] = 0;
}

template<typename idx_t, typename data_t, typename calc_t, int dof>
CSRMatrix<idx_t, data_t, calc_t, dof>::CSRMatrix(const idx_t _nrows, const idx_t _ncols, const idx_t _nnz): nrows(_nrows), ncols(_ncols), nnz(_nnz)
{
    alloc_mem();
}

template<typename idx_t, typename data_t, typename calc_t, int dof>
CSRMatrix<idx_t, data_t, calc_t, dof>::~CSRMatrix()
{
    delete [] row_ptr;
    delete [] col_idx;
    delete [] vals;
}

template<typename idx_t, typename data_t, typename calc_t, int dof>
void CSRMatrix<idx_t, data_t, calc_t, dof>::alloc_mem(bool zero_rpt)
{
    if (row_ptr) delete [] row_ptr;
    if (col_idx) delete [] col_idx;
    if (vals   ) delete [] vals;
    
    row_ptr = new idx_t [nrows + 1];
    col_idx = new idx_t [nnz];
    const long long tot_len = ((long long) nnz) * dof * dof;
    vals    = new data_t[tot_len];

    if (zero_rpt) {
        #pragma omp parallel for schedule(static)
        for (idx_t i = 0; i < nrows + 1; i++)
            row_ptr[i] = 0;
    }
}

template<typename idx_t, typename data_t, typename calc_t, int dof>
void CSRMatrix<idx_t, data_t, calc_t, dof>::Mult(calc_t alpha, const calc_t *x, calc_t beta, const calc_t * b, calc_t *y) const
{
    const bool nonexist_b = (b == nullptr || beta == 0.0);
    if constexpr (dof == 1) {
        if (nonexist_b) {
            #pragma omp parallel for schedule(static)
            for (idx_t i = 0; i < nrows; i++) {
                calc_t tmp = 0.0;
                for (idx_t p = row_ptr[i]; p < row_ptr[i+1]; p++) {
                    tmp += vals[p] * x[col_idx[p]];
                }
                y[i] = alpha * tmp;
            }
        } else {
            #pragma omp parallel for schedule(static)
            for (idx_t i = 0; i < nrows; i++) {
                calc_t tmp = 0.0;
                for (idx_t p = row_ptr[i]; p < row_ptr[i+1]; p++) {
                    tmp += vals[p] * x[col_idx[p]];
                }
                y[i] = alpha * tmp + beta * b[i];
            }
        }
    } else {
        constexpr int e_size = dof*dof;
        if (nonexist_b) {
            #pragma omp parallel for schedule(static)
            for (idx_t i = 0; i < nrows; i++) {
                calc_t tmp[dof]; d_vec_zero<calc_t, dof>(tmp);
                for (idx_t p = row_ptr[i]; p < row_ptr[i+1]; p++) {
                    const idx_t j = col_idx[p];
                    matvec_mla<idx_t, data_t, calc_t, dof>(vals + p * e_size, x + j * dof, tmp); 
                }
                calc_t * dst_ptr = y + i * dof;
                #pragma GCC unroll (4) 
                for (int f = 0; f < dof; f++) dst_ptr[f] = alpha * tmp[f];// y[i] = alpha * tmp;
            }
        } else {
            #pragma omp parallel for schedule(static)
            for (idx_t i = 0; i < nrows; i++) {
                calc_t tmp[dof]; d_vec_zero<calc_t, dof>(tmp);
                for (idx_t p = row_ptr[i]; p < row_ptr[i+1]; p++) {
                    const idx_t j = col_idx[p];
                    matvec_mla<idx_t, data_t, calc_t, dof>(vals + p * e_size, x + j * dof, tmp); 
                }
                calc_t       * dst_ptr = y + i * dof;
                const calc_t * src_ptr = b + i * dof;
                #pragma GCC unroll (4) 
                for (int f = 0; f < dof; f++) dst_ptr[f] = alpha * tmp[f] + beta * src_ptr[f];
            }
        }
    }
}

template<typename idx_t, typename data_t, typename calc_t, int dof>
void CSRMatrix<idx_t, data_t, calc_t, dof>::transpose(CSRMatrix<idx_t, data_t, calc_t, dof> & myT)
{
    CSRMatrixTranspose<idx_t, data_t, calc_t, data_t, calc_t, dof>(*this, myT);
}

template class CSRMatrix<int, float , float , 1>;
template class CSRMatrix<int, float , float , 2>;
template class CSRMatrix<int, float , float , 3>;
template class CSRMatrix<int, float , float , 4>;

template class CSRMatrix<int, float , double, 1>;
template class CSRMatrix<int, float , double, 2>;
template class CSRMatrix<int, float , double, 3>;
template class CSRMatrix<int, float , double, 4>;

template class CSRMatrix<int, double, double, 1>;
template class CSRMatrix<int, double, double, 2>;
template class CSRMatrix<int, double, double, 3>;
template class CSRMatrix<int, double, double, 4>;
#ifdef USE_FP16
template class CSRMatrix<int, __fp16, float , 1>;
template class CSRMatrix<int, __fp16, float , 2>;
template class CSRMatrix<int, __fp16, float , 3>;
template class CSRMatrix<int, __fp16, float , 4>;

template class CSRMatrix<int, __fp16, double, 1>;
template class CSRMatrix<int, __fp16, double, 2>;
template class CSRMatrix<int, __fp16, double, 3>;
template class CSRMatrix<int, __fp16, double, 4>;
#endif

